from llm import llm
from llm import super_eval
from prompt import qa3
import os
import hashlib


class QA:
    def __init__(self):
        self.system_prompt = qa3
        self.qa_path = []

    def add_qa(self, q, a):
        self.qa_path.append({'q': q, 'a': a})

    def add_q(self, q):
        self.qa_path.append({'q': q})

    def add_a(self, a):
        self.qa_path.append({'a': a})

    def get_new_question(self, q):
        messages = [{'role': 'system', 'content': self.system_prompt}]
        user_content = ''
        n = 0
        for i in self.qa_path:
            if 'q' in i:
                n += 1
            if 'q' in i:
                user_content += f"Q{n}:{i['q']}\n"
            if 'a' in i:
                user_content += f"A{n}:{i['a']}\n"
        n += 1
        user_content += f"Q{n}:{q}\n"
        messages.append({'role': 'user',
                         'content': user_content + '\n请输出json:'})
        res = super_eval(llm(messages))

        return res


def process_answer(messages):
    if len(messages) < 6:
        return []
    check = {}
    check['in'] = messages

    if len(messages) == 6:
        messages[-2]['content'] = '内容略'
        if '查询结果: []' in messages[-2]['content']:
            print('已省略代码')
            messages[-3]['content'] = '代码略'
            return messages  # [:-2]
        else:
            return messages
    messages[-2]['content'] = '内容略'

    prompt = '''
请按照最终结果，输出回答内容，即
--start--
分析字段
xxx
选择字段
```python
doc_searcher.choose([{"table_name": "xxx", "columns": ["xxx",...]}]) # 这里是所需的所有字段
```
构造查询代码
```python
这里是解决问题的完整代码，请按照之前的对话
```
请执行以上代码获取结果。
--end--
这部分内容，按照最终结果套用上面的格式重新输出，仅仅输出内容：
--start--
'''
    messages.append({'role': 'user', 'content': prompt})

    answer = llm(messages)

    res = messages[:3] + [{'role': 'assistant', 'content': answer}]

    check['out'] = res
    with open('check.json', 'a', encoding='utf8') as f:
        f.write(json.dumps(check, ensure_ascii=False))
        f.write('\n')
    return res


if __name__ == '__main__':
    import json

    n = 0
    with open("sample_orm.json", encoding='utf8') as f:
        for i in f:
            data = json.loads(i)
            if len(data) > 6:
                process_answer(data)
            # print(f'====={len(data)}======')

    # question =     {'tid': 'tttt----7',
    #  'team': [
    #      {'id': 'tttt----7----7-2-1', 'question': '天顺风能属于哪个三级行业？', 'answer': 'ThirdIndustryName 风电零部件',
    #       'sql': "\nSELECT ThirdIndustryName\nFROM AStockIndustryDB.LC_ExgIndustry\nWHERE CompanyCode = 81722\nAND IfPerformed = 1\nORDER BY DATE(STR_TO_DATE(InfoPublDate, '%Y-%m-%d %H:%i:%s.%f')) DESC\nLIMIT 1;\n"},
    #      {'id': 'tttt----7----7-2-2', 'question': '2021年发布的该行业的股票有多少只？',
    #       'answer': '2021年发布的该行业的股票有14只。',
    #       'sql': "\nSELECT COUNT(DISTINCT CompanyCode) AS StockCount\nFROM AStockIndustryDB.LC_ExgIndustry\nWHERE ThirdIndustryName = '风电零部件'\nAND IfPerformed = 1\nAND YEAR(STR_TO_DATE(InfoPublDate, '%Y-%m-%d %H:%i:%s.%f')) = 2021;\n"},
    #      {'id': 'tttt----7----7-2-3', 'question': '该行业8月公布的最高行业营业收入是多少万元？',
    #       'answer': '该行业8月公布的最高行业营业收入是1571591.3774万元。',
    #       'sql': "\nSELECT MAX(IndOperatingRevenue) AS MaxOperatingRevenue\nFROM AStockIndustryDB.LC_IndFinIndicators\nWHERE IndustryName = '风电零部件'\nAND MONTH(STR_TO_DATE(InfoPublDate, '%Y-%m-%d %H:%i:%s.%f')) = 8\nAND YEAR(STR_TO_DATE(InfoPublDate, '%Y-%m-%d %H:%i:%s.%f')) = 2021;\n"}]}
    #
    # qa  = QA()
    #
    # for inx,i in enumerate(question['team']):
    #     qa.add_qa(i['question'],i['answer'])
    #     r = qa.get_new_question(question['team'][inx+1]['question'])
    #     print(r)
